import os
import sys
import yxn
import time
import json
import socket
import base64
import requests
from log import logger
from datetime import datetime, timedelta

# 开始结束时间
START_TS = int(sys.argv[1])
END_TS   = int(sys.argv[2])

# LOGSTASH地址
LOGSTASH_ADDR = "172.31.0.13"
LOGSTASH_PORT = 8881

# 地址、采样率等参数
UNISERVER_ADDR = "172.31.0.13"
# web端口、数据端口、用户名密码
UNISERVER_WEB_PORT = 9000
UNISERVER_DATA_PORT = 12201
UNISERVER_USERNAME="admin"
UNISERVER_PASSWORD="Synmesh@2023"
# 前端负载均衡采样率
UNISERVER_SAMPLE = 1
# 聚合结果发送协议
UNISERVER_SEND_PROTOCOL = "http"
# 仪表板ID和查询ID，根据浏览器审查元素获取
UNISERVER_DASHBOARD_ID="6721e50cef72b4c0274e65a1"

# 主机维度结果集ID
UNISERVER_HOST_QUERY_ID="44f56288-759a-49b9-9854-64ee11ade27a"
UNISERVER_HOST_QUERY_GROUP_TITLE_LIST=["source_collector", "source", "host", "message", "SYSTEM_NAME", "HOST_SYSTEM", "SOURCE"]
UNISERVER_HOST_QUERY_ALG_TITLE_LIST=["HOST_BYTES", "HOST_NET_DELAY_AVG", "HOST_NET_DELAY_MAX", "HOST_CPU_CORE", "HOST_MEM_TOTAL", "HOST_DISK_TOTAL", "HOST_CPU_AVG", "HOST_MEM_AVG", "HOST_DISK_AVG", "HOST_CPU_MAX", "HOST_MEM_MAX", "HOST_DISK_MAX", "HOST_CPU_CORE_USE"]
UNISERVER_HOST_QUERY_PRIKEY_IDX_LIST=[0]
# 模块维度结果集ID
UNISERVER_MODULE_QUERY_ID="4dfcea9e-9307-4640-b5c6-e7972bd99d33"
UNISERVER_MODULE_QUERY_GROUP_TITLE_LIST=["MODULE", "APP_NAME", "APP_ID"]
UNISERVER_MODULE_QUERY_ALG_TITLE_LIST=["APP_HTTP_500_COUNT", "APP_HTTP_400_COUNT", "APP_HTTP_SLOW_COUNT", "APP_HTTP_SLOW_AVG", "APP_HTTP_SLOW_MAX", "APP_BYTES", "APP_HTTP_RESPONSE_MAX", "APP_HTTP_RESPONSE_AVG", "APP_HOME_PAGE_COUNT", "APP_HTTP_COUNT"]
UNISERVER_MODULE_QUERY_PRIKEY_IDX_LIST=[0, 1]
# 应用维度结果集ID
UNISERVER_APP_QUERY_ID="7297ff8d-77fa-4815-a3dd-18243e9f9964"
UNISERVER_APP_QUERY_GROUP_TITLE_LIST=["APP_NAME", "APP_ID"]
UNISERVER_APP_QUERY_ALG_TITLE_LIST=["APP_HTTP_500_COUNT", "APP_HTTP_400_COUNT", "APP_HTTP_SLOW_COUNT", "APP_HTTP_SLOW_AVG", "APP_HTTP_SLOW_MAX", "APP_BYTES", "APP_HTTP_RESPONSE_MAX", "APP_HTTP_RESPONSE_AVG", "APP_HOME_PAGE_COUNT", "APP_HTTP_COUNT"]
UNISERVER_APP_QUERY_PRIKEY_IDX_LIST=[0]
# 系统维度结果集ID
UNISERVER_SYSTEM_QUERY_ID="bbec800e-405f-4dac-abf8-e483a129cf0f"
UNISERVER_SYSTEM_QUERY_GROUP_TITLE_LIST=["SYSTEM", "HOST_SYSTEM", "SOURCE", "SOURCE_ID", "source_collector"]
UNISERVER_SYSTEM_QUERY_ALG_TITLE_LIST=["SYSTEM_CPU_CORE_TOTAL", "SYSTEM_MEM_TOTAL", "SYSTEM_DISK_TOTAL", "SYSTEM_MEM_USAGE", "SYSTEM_DISK_USAGE", "SYSTEM_MEM", "SYSTEM_DISK"]
UNISERVER_SYSTEM_QUERY_PRIKEY_IDX_LIST=[0]
# 网络维度结果集ID
UNISERVER_NETWORK_QUERY_ID="d8a28db7-5c82-433a-a494-6735fa71be8b"
UNISERVER_NETWORK_QUERY_GROUP_TITLE_LIST=["NETWORK", "APP_NAME", "APP_ID"]
UNISERVER_NETWORK_QUERY_ALG_TITLE_LIST=["APP_BYTES", "APP_HOME_PAGE_COUNT", "APP_HTTP_COUNT", "APP_HTTP_500_COUNT", "APP_HTTP_400_COUNT", "APP_HTTP_SLOW_COUNT"]
UNISERVER_NETWORK_QUERY_PRIKEY_IDX_LIST=[0, 1]
# 单位维度结果集ID
UNISERVER_DEPARTMENT_QUERY_ID="a394f929-4cb0-413f-8fd7-9ab516ca2b75"
UNISERVER_DEPARTMENT_QUERY_GROUP_TITLE_LIST=["SOURCE", "SOURCE_ID", "source_collector"]
UNISERVER_DEPARTMENT_QUERY_ALG_TITLE_LIST=["DEPARTMENT_CPU_CORE_TOTAL", "DEPARTMENT_MEM_TOTAL", "DEPARTMENT_DISK_TOTAL", "DEPARTMENT_MEM_USAGE", "DEPARTMENT_DISK_USAGE", "DEPARTMENT_MEM", "DEPARTMENT_DISK"]
UNISERVER_DEPARTMENT_QUERY_PRIKEY_IDX_LIST=[0]

# 仪表板元素数据
METADATA_FILE = "metadata.json"
# 登录请求session地址
TOKEN_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/system/sessions"
# 提交仪表板元数据地址
METADATA_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/views/search/metadata"
# 执行查询获取结果地址
EXECUTE_URL = f"http://{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}/api/views/search/{UNISERVER_DASHBOARD_ID}/execute"

# 聚合结果文件保存路径
HOST_ARCH_RESULT_SAVE_PATH = "/uniarch/host"
MODULE_ARCH_RESULT_SAVE_PATH = "/uniarch/module"
APP_ARCH_RESULT_SAVE_PATH = "/uniarch/app"
SYSTEM_ARCH_RESULT_SAVE_PATH = "/uniarch/system"
NETWORK_ARCH_RESULT_SAVE_PATH = "/uniarch/network"
DEPARTMENT_ARCH_RESULT_SAVE_PATH = "/uniarch/department"

# 模拟登录请求体
TOKEN_DATA = {
  "username": UNISERVER_USERNAME,
  "password": UNISERVER_PASSWORD,
  "host": f"{UNISERVER_ADDR}:{UNISERVER_WEB_PORT}"
}

# 模拟查询仪表板请求体
EXECUTE_DATA = {
  "global_override": {
    "timerange": {
      "type": "absolute",
      "from": None,
      "to": None
    }
  },
  "parameter_bindings": {}
}

# 模拟任何请求的头
REQUEST_HEADERS = {
  "Accept": "application/json",
  "Accept-Language": "zh-CN,zh;q=0.9",
  "Content-Type": "application/json",
  "X-Requested-By": "XMLHttpRequest",
  "X-Requested-With": "XMLHttpRequest"
}

def send_arch_merge_result(agg_result):
  # 发送到GELF HTTP 或者 UDP
  if 'http' == UNISERVER_SEND_PROTOCOL:
    gelf_url = f'http://{UNISERVER_ADDR}:{UNISERVER_DATA_PORT}/gelf'
    with requests.Session() as session:
      for key, val in agg_result.items():
        message = json.dumps(val)
        session.post(gelf_url, data=message, headers={'Content-Type': 'application/json'})
  elif 'udp' == UNISERVER_SEND_PROTOCOL:
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as gelf_udp:
      for key, val in agg_result.items():
        message = json.dumps(val)
        gelf_udp.sendto(message.encode('utf-8'), (UNISERVER_ADDR, UNISERVER_DATA_PORT))
  with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as gelf_udp:
    for key, val in agg_result.items():
      message = json.dumps(val)
      gelf_udp.sendto(message.encode('utf-8'), (LOGSTASH_ADDR, LOGSTASH_PORT))
      time.sleep(0.05)

def save_arch_merge_result(dir_path, agg_result):
  # 目录不存在，创建目录
  if not os.path.exists(dir_path): os.makedirs(dir_path)
  # 保存到文件存储
  filename = datetime.fromtimestamp(START_TS).strftime("%Y-%m-%d %H:%M:%S")
  result_file = f'{dir_path}/{filename}.json'
  with open(result_file, 'w', encoding='utf-8') as file:
    json.dump(agg_result, file, ensure_ascii=False, indent=2)
  # 打印保存主机数量
  logger.info(f'MERGE: save {len(agg_result)} records to {result_file} ...')

# 处理仪表板每个小部件中的一行
# 'key': ['734599345345345', '10.0.8.2'],
# 'values': [{
#   'key': ['HOST_SESSION_COUNT'],
#   'value': 29406,
#   'rollup': True,
#   'source': 'row-leaf'
# }, {
def process_metrics_row(group_prikey_idx_list, group_title_list, alg_title_list, metrics_row, agg_result, arch_type):
  # 获取分组部分和算法部分
  group_field_list = metrics_row['key']
  alg_field_list = metrics_row['values']

  # 分组字段不等于分组标题数量
  if len(group_title_list) != len(group_field_list): return

  prikey = '' # 拼接主键
  for group_prikey_idx in group_prikey_idx_list:
    prikey += group_field_list[group_prikey_idx]

  # 如果不存在创建主键字典
  if prikey not in agg_result:
    agg_result[prikey] = yxn.create_metrics_dict(group_title_list, alg_title_list, group_field_list, START_TS, END_TS, arch_type)

  # 更新主键归档行分组拼接字段集合，用于统计系统、部门下主机数量
  group_field_list_str = '_'.join(group_field_list)
  agg_result[prikey]['ARCH_SET'].add(group_field_list_str)

  # 对算法指标进行赋值或累加
  for alg_dict in alg_field_list:
    key = alg_dict['key'][0]
    val = alg_dict['value']
    # 过滤掉值为空（-Infinity）的算法结果
    if isinstance(val, str): continue
    # 有的数据里内存和磁盘总量用字节计算，单独处理
    if 'MEM_TOTAL' in key or 'DISK_TOTAL' in key: 
      if val > 1073741824: val = int(val / 1073741824)
    # 对key进行赋值，重复出现此key累加
    metrics_dict = agg_result[prikey]
    metrics_dict[key] += val

# 输入仪表板查询结果，返回聚合结果
def process_query_result(group_prikey_idx_list, group_title_list, alg_title_list, query_result, arch_type):
  agg_result = {}
  # 获取仪表板查询结果，遍历仪表板的小部件
  for unit_id, unit_result in query_result.items():
    # 小部件包括多个行，每行有很多指标
    metrics_row_list = unit_result['rows']
    # 遍历每行的指标字典
    for metrics_row in metrics_row_list:
      process_metrics_row(group_prikey_idx_list, group_title_list, alg_title_list, metrics_row, agg_result, arch_type)
  # 取归档行分组拼接字段集合长度，用于统计系统、部门下主机数量
  for key, metrics_dict in agg_result.items():
    metrics_dict['ARCH_COUNT'] = len(metrics_dict['ARCH_SET'])
    metrics_dict.pop('ARCH_SET')
  return agg_result

# 模拟登录查询仪表板
def query_dashboard():
  # 读取仪表板查询JSON
  metadata_data = None
  with open(METADATA_FILE, "r") as f: metadata_data = json.load(f)
  
  # 设置查询时间
  EXECUTE_DATA['global_override']['timerange']['from'] = datetime.utcfromtimestamp(START_TS).isoformat() + "Z"
  EXECUTE_DATA['global_override']['timerange']['to'] = datetime.utcfromtimestamp(END_TS).isoformat() + "Z"
  logger.info(EXECUTE_DATA)
  
  # 模拟登录获取session
  logger.info(f'get token from {TOKEN_URL} ...')
  response = requests.post(TOKEN_URL, headers=REQUEST_HEADERS, json=TOKEN_DATA, verify=False)
  logger.info(f'token response {response.text} ...')
  session_id = json.loads(response.text).get('session_id', None)
  
  # 根据session_id生成token令牌
  uniserver_token = base64.b64encode(f'{session_id}:session'.encode('utf-8')).decode('utf-8')
  REQUEST_HEADERS["Authorization"] = f"Basic {uniserver_token}"
  logger.info(f'token is: {uniserver_token} ...')
  
  # 查询仪表板
  logger.info(f'query {START_TS} to {END_TS} ...')
  requests.post(METADATA_URL, headers=REQUEST_HEADERS, json=metadata_data, verify=False)
  response = requests.post(EXECUTE_URL, headers=REQUEST_HEADERS, json=EXECUTE_DATA, verify=False)
  
  # 解析查询结果
  query_dashboard_result = None
  try:
    query_dashboard_result = json.loads(response.text)
    if "results" not in query_dashboard_result:
      logger.warning(f'no query results in {response.text}')
  except:
    logger.warning(f'failed to load {response.text}')

  return query_dashboard_result

# 查询仪表板
query_dashboard_result = query_dashboard()

# 处理主机仪表板查询结果
host_query_result = query_dashboard_result["results"][UNISERVER_HOST_QUERY_ID]["search_types"]
host_agg_result = process_query_result(UNISERVER_HOST_QUERY_PRIKEY_IDX_LIST, 
                                       UNISERVER_HOST_QUERY_GROUP_TITLE_LIST, 
                                       UNISERVER_HOST_QUERY_ALG_TITLE_LIST, 
                                       host_query_result, 
                                       "HOST")
save_arch_merge_result(HOST_ARCH_RESULT_SAVE_PATH, host_agg_result)
send_arch_merge_result(host_agg_result)
 
# 处理模块仪表板查询结果
module_query_result = query_dashboard_result["results"][UNISERVER_MODULE_QUERY_ID]["search_types"]
module_agg_result = process_query_result(UNISERVER_MODULE_QUERY_PRIKEY_IDX_LIST, 
                                         UNISERVER_MODULE_QUERY_GROUP_TITLE_LIST, 
                                         UNISERVER_MODULE_QUERY_ALG_TITLE_LIST, 
                                         module_query_result, 
                                         "MODULE")
save_arch_merge_result(MODULE_ARCH_RESULT_SAVE_PATH, module_agg_result)
send_arch_merge_result(module_agg_result)
 
# 处理应用仪表板查询结果
app_query_result = query_dashboard_result["results"][UNISERVER_APP_QUERY_ID]["search_types"]
app_agg_result = process_query_result(UNISERVER_APP_QUERY_PRIKEY_IDX_LIST,
                                      UNISERVER_APP_QUERY_GROUP_TITLE_LIST,
                                      UNISERVER_APP_QUERY_ALG_TITLE_LIST,
                                      app_query_result, 
                                      "APP")
save_arch_merge_result(APP_ARCH_RESULT_SAVE_PATH, app_agg_result)
send_arch_merge_result(app_agg_result)

# 处理系统仪表板查询结果
system_query_result = query_dashboard_result["results"][UNISERVER_SYSTEM_QUERY_ID]["search_types"]
system_agg_result = process_query_result(UNISERVER_SYSTEM_QUERY_PRIKEY_IDX_LIST, 
                                         UNISERVER_SYSTEM_QUERY_GROUP_TITLE_LIST,
                                         UNISERVER_SYSTEM_QUERY_ALG_TITLE_LIST,
                                         system_query_result, 
                                         "SYSTEM")
for name, metrics_dict in system_agg_result.items():
  if not yxn.calc_metrics_ratio(metrics_dict, "SYSTEM_MEM_USAGE", "SYSTEM_MEM_TOTAL", "SYSTEM_MEM", True):
    logger.warning(f'{name} not have mem ratio')
  if not yxn.calc_metrics_ratio(metrics_dict, "SYSTEM_DISK_USAGE", "SYSTEM_DISK_TOTAL", "SYSTEM_DISK", True):
    logger.warning(f'{name} not have disk ratio')
save_arch_merge_result(SYSTEM_ARCH_RESULT_SAVE_PATH, system_agg_result)
send_arch_merge_result(system_agg_result)

# 处理网络仪表板查询结果
network_query_result = query_dashboard_result["results"][UNISERVER_NETWORK_QUERY_ID]["search_types"]
network_agg_result = process_query_result(UNISERVER_NETWORK_QUERY_PRIKEY_IDX_LIST, 
                                          UNISERVER_NETWORK_QUERY_GROUP_TITLE_LIST,
                                          UNISERVER_NETWORK_QUERY_ALG_TITLE_LIST,
                                          network_query_result, 
                                          "NETWORK")
save_arch_merge_result(NETWORK_ARCH_RESULT_SAVE_PATH, network_agg_result)
send_arch_merge_result(network_agg_result)

# 处理单位仪表板查询结果
department_query_result = query_dashboard_result["results"][UNISERVER_DEPARTMENT_QUERY_ID]["search_types"]
department_agg_result = process_query_result(UNISERVER_DEPARTMENT_QUERY_PRIKEY_IDX_LIST, 
                                            UNISERVER_DEPARTMENT_QUERY_GROUP_TITLE_LIST,
                                            UNISERVER_DEPARTMENT_QUERY_ALG_TITLE_LIST,
                                            department_query_result, 
                                            "DEPARTMENT")
for name, metrics_dict in department_agg_result.items():
  if not yxn.calc_metrics_ratio(metrics_dict, "DEPARTMENT_MEM_USAGE", "DEPARTMENT_MEM_TOTAL", "DEPARTMENT_MEM", True):
    logger.warning(f'{name} not have mem ratio')
  if not yxn.calc_metrics_ratio(metrics_dict, "DEPARTMENT_DISK_USAGE", "DEPARTMENT_DISK_TOTAL", "DEPARTMENT_DISK", True):
    logger.warning(f'{name} not have disk ratio')
save_arch_merge_result(DEPARTMENT_ARCH_RESULT_SAVE_PATH, department_agg_result)
send_arch_merge_result(department_agg_result)

